# performance_profiler.py (修正版本)

import os
import sys
import time
import cProfile
import pstats
from io import StringIO

# 修正路径设置 - 获取项目根目录
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(current_dir))  # 向上两级到 RAG 目录
sys.path.insert(0, project_root)

from knowledge.core.kb_application import KnowledgeBaseApplication

class PerformanceProfiler:
    """性能分析器"""
    
    def __init__(self):
        self.kb_app = KnowledgeBaseApplication("enterprise_knowledge")
        self.kb_app.initialize()
        # 使用现有的方法，不调用 ensure_query_engine_loaded
        self._load_query_engine()
    
    def _load_query_engine(self):
        """加载查询引擎"""
        try:
            # 通过构建知识库来加载查询引擎
            knowledge_root = "./resource"
            if os.path.exists(knowledge_root):
                # 这会加载现有的索引
                stats = self.kb_app.build_knowledge_base(knowledge_root, force_rebuild=False)
                if stats.get('total_documents', 0) > 0:
                    print("✅ 查询引擎加载成功")
                else:
                    print("❌ 查询引擎加载失败")
            else:
                print("❌ 知识库目录不存在")
        except Exception as e:
            print(f"❌ 加载查询引擎失败: {e}")
    
    def profile_query_components(self, query_text):
        """分析查询各个组件的性能"""
        print(f"\n🔍 性能分析: '{query_text}'")
        print("=" * 80)
        
        total_start = time.time()
        
        # 检查查询引擎是否可用
        if not self.kb_app.query_engine:
            print("❌ 查询引擎未初始化")
            return None
        
        # 1. 检索阶段性能
        retrieval_start = time.time()
        try:
            retriever = self.kb_app.query_engine.query_engine.retriever
            retrieval_results = retriever.retrieve(query_text)
            retrieval_time = time.time() - retrieval_start
            
            print(f"📊 检索阶段:")
            print(f"   ⏱️  耗时: {retrieval_time:.3f}秒")
            print(f"   📄 检索到 {len(retrieval_results)} 个文档片段")
            
            # 显示检索结果的质量
            if retrieval_results:
                print(f"   🎯 最佳匹配分数: {retrieval_results[0].score:.4f}")
                for i, result in enumerate(retrieval_results[:3], 1):
                    print(f"      {i}. 分数: {result.score:.4f}, 文件: {result.node.metadata.get('file_name', 'N/A')}")
        except Exception as e:
            print(f"   ❌ 检索失败: {e}")
            retrieval_time = 0
            retrieval_results = []
        
        # 2. 完整查询性能
        full_start = time.time()
        try:
            # 使用标准查询方法
            full_result = self.kb_app.query(query_text, show_details=False)
            full_time = time.time() - full_start
            
            print(f"🎯 完整查询:")
            print(f"   ⏱️  总耗时: {full_time:.3f}秒")
            if retrieval_time > 0:
                print(f"   📈 检索耗时占比: {retrieval_time/full_time*100:.1f}%")
                print(f"   🤖 LLM生成耗时: {full_time - retrieval_time:.3f}秒")
                print(f"   📊 LLM占比: {(full_time - retrieval_time)/full_time*100:.1f}%")
            
            # 显示查询结果摘要
            self._display_query_result(full_result)
            
            return full_result
        except Exception as e:
            print(f"   ❌ 完整查询失败: {e}")
            return None
    
    def _display_query_result(self, result):
        """显示查询结果摘要"""
        if not result:
            print("   ❌ 无查询结果")
            return
        
        # 根据 kb_application.py 中的查询结果结构处理
        if isinstance(result, dict):
            # 标准查询结果结构
            if 'answer' in result:
                answer = result['answer']
                response_length = len(str(answer))
                print(f"   📝 响应长度: {response_length} 字符")
                print(f"   💬 回答摘要: {answer[:100]}..." if len(answer) > 100 else f"   💬 回答: {answer}")
            
            elif 'response' in result:
                # 兼容旧版本结构
                response = result['response']
                response_length = len(str(response))
                print(f"   📝 响应长度: {response_length} 字符")
                print(f"   💬 回答摘要: {response[:100]}..." if len(response) > 100 else f"   💬 回答: {response}")
            
            # 显示检索结果数量
            if 'retrieval_results' in result:
                print(f"   📄 检索片段: {len(result['retrieval_results'])} 个")
            elif 'sources' in result:
                print(f"   📄 来源文档: {len(result['sources'])} 个")
                
        elif isinstance(result, str):
            # 直接返回字符串的情况
            response_length = len(result)
            print(f"   📝 响应长度: {response_length} 字符")
            print(f"   💬 回答摘要: {result[:100]}..." if len(result) > 100 else f"   💬 回答: {result}")
    
    def compare_queries(self, queries):
        """比较多个查询的性能"""
        print("\n" + "="*80)
        print("📊 多查询性能对比")
        print("="*80)
        
        results = []
        for query in queries:
            start_time = time.time()
            try:
                result = self.kb_app.query(query, show_details=False)
                end_time = time.time()
                
                response_time = end_time - start_time
                
                # 根据实际结果结构计算响应长度
                response_length = self._get_response_length(result)
                retrieval_count = self._get_retrieval_count(result)
                
                results.append({
                    'query': query,
                    'time': response_time,
                    'response_length': response_length,
                    'retrieval_count': retrieval_count,
                    'result': result
                })
                
                print(f"\n🔍 '{query}'")
                print(f"   ⏱️  耗时: {response_time:.3f}秒")
                print(f"   📝 响应长度: {response_length} 字符")
                print(f"   📄 检索片段: {retrieval_count} 个")
                
            except Exception as e:
                print(f"\n🔍 '{query}' - 失败: {e}")
                results.append({
                    'query': query,
                    'time': 0,
                    'response_length': 0,
                    'retrieval_count': 0,
                    'error': str(e)
                })
        
        # 分析性能差异
        if results:
            print(f"\n📈 性能分析:")
            valid_results = [r for r in results if r['time'] > 0]
            if valid_results:
                avg_time = sum(r['time'] for r in valid_results) / len(valid_results)
                print(f"   平均耗时: {avg_time:.3f}秒")
                
                for r in valid_results:
                    deviation = (r['time'] - avg_time) / avg_time * 100
                    print(f"   '{r['query'][:20]}...': {r['time']:.3f}s ({deviation:+.1f}%)")
    
    def _get_response_length(self, result):
        """从查询结果中获取响应长度"""
        if not result:
            return 0
        
        if isinstance(result, dict):
            if 'answer' in result:
                return len(str(result['answer']))
            elif 'response' in result:
                return len(str(result['response']))
        elif isinstance(result, str):
            return len(result)
        
        return len(str(result))
    
    def _get_retrieval_count(self, result):
        """从查询结果中获取检索片段数量"""
        if not result or not isinstance(result, dict):
            return 0
        
        if 'retrieval_results' in result:
            return len(result['retrieval_results'])
        elif 'sources' in result:
            return len(result['sources'])
        elif 'source_nodes' in result:
            return len(result['source_nodes'])
        
        return 0
    
    def test_retrieval_only(self, query_text):
        """只测试检索性能，不调用LLM"""
        print(f"\n🧪 纯检索测试: '{query_text}'")
        print("=" * 60)
        
        if not self.kb_app.query_engine:
            print("❌ 查询引擎未初始化")
            return
        
        start_time = time.time()
        try:
            retriever = self.kb_app.query_engine.query_engine.retriever
            retrieval_results = retriever.retrieve(query_text)
            end_time = time.time()
            
            print(f"⏱️  纯检索耗时: {end_time - start_time:.3f}秒")
            print(f"📄 检索到 {len(retrieval_results)} 个文档片段")
            
            for i, result in enumerate(retrieval_results[:3], 1):
                metadata = result.node.metadata
                content_preview = result.node.text[:100] + "..." if len(result.node.text) > 100 else result.node.text
                print(f"{i}. [{result.score:.4f}] {metadata.get('file_name', 'N/A')}")
                print(f"   内容: {content_preview}")
                print()
                
        except Exception as e:
            print(f"❌ 检索测试失败: {e}")

def main():
    """主函数"""
    print("🚀 开始性能分析...")
    
    try:
        profiler = PerformanceProfiler()
        
        # 测试用例
        slow_query = "你知道JRT公司吗？"
        fast_query = "JRT公司2024年的研发投入和收入是多少？"
        
        # 1. 分析慢查询
        print("\n" + "="*80)
        profiler.profile_query_components(slow_query)
        
        # 2. 分析快查询
        print("\n" + "="*80)
        profiler.profile_query_components(fast_query)
        
        # 3. 纯检索性能测试
        print("\n" + "="*80)
        profiler.test_retrieval_only(slow_query)
        profiler.test_retrieval_only(fast_query)
        
        # 4. 对比多个查询
        print("\n" + "="*80)
        test_queries = [
            "你知道JRT公司吗？",
            "JRT公司2024年的研发投入和收入是多少？", 
            "差旅报销标准是什么？",
            "北京出差的住宿标准",
            "JRT公司的市场份额"
        ]
        
        profiler.compare_queries(test_queries)
        
    except Exception as e:
        print(f"❌ 性能分析失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()